#include "FrameBuffer.h"

FrameBuffer::FrameBuffer()
{
    mFBO = 0;
    mDepthBuffer = nullptr;
}

FrameBuffer::~FrameBuffer()
{
    for (auto it = mAttachments.begin(); it != mAttachments.end(); ++it)
    {
        delete *it;
    }

    if (mFBO)
    {
        vkDestroyFramebuffer(xGetVulkanDevice(), mFBO, nullptr);
    }

    if (mRenderPass)
    {
        vkDestroyRenderPass(xGetVulkanDevice(), mRenderPass, nullptr);
    }
}

void FrameBuffer::SetSize(uint32_t width, uint32_t height)
{
    mWidth = width;
    mHeight = height;
}

void FrameBuffer::AttachColorBuffer(VkFormat format)
{
    Texture2D *colorBuffer = new Texture2D();
    colorBuffer->mFormat = format;

    xGenImage(colorBuffer, mWidth, mHeight, format, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT);
    xGenImageView2D(colorBuffer);
    xGenSampler(colorBuffer);

    ++mColorBufferCount;
    mAttachments.push_back(colorBuffer);

    VkClearValue cv;
    cv.color = {0.0f, 0.f, 0.0f, 0.0f};
    mClearValues.push_back(cv);

    mDepthBufferIndex = mAttachments.size();
}

void FrameBuffer::AttachDepthBuffer()
{
    Texture2D *depth_buffer = new Texture2D(VK_IMAGE_ASPECT_DEPTH_BIT);
    depth_buffer->mFormat = VK_FORMAT_D24_UNORM_S8_UINT;

    xGenImage(depth_buffer, mWidth, mHeight, depth_buffer->mFormat, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT);
    xGenImageView2D(depth_buffer);
    xGenSampler(depth_buffer);

    mAttachments.push_back(depth_buffer);
    mDepthBuffer = depth_buffer;

    VkClearValue cv;
    cv.depthStencil = {1.0f, 0};
    mClearValues.push_back(cv);
}

void FrameBuffer::Finish()
{
    std::vector<VkAttachmentDescription> attachments{};
    attachments.resize(mAttachments.size());

    std::vector<VkAttachmentReference> color_attachment_refs{};
    VkAttachmentReference depthAttachmentRef{};

    std::vector<VkImageView> render_targets{};
    render_targets.resize(mAttachments.size());

    int color_buffer_count = 0;
    for (size_t i = 0; i < mAttachments.size(); ++i)
    {
        Texture2D *texture = mAttachments[i];
        if (texture->mImageAspectFlag == VK_IMAGE_ASPECT_COLOR_BIT)
        {
            ++color_buffer_count;
        }
        else if (texture->mImageAspectFlag == VK_IMAGE_ASPECT_DEPTH_BIT)
        {
        }
        render_targets[i] = texture->mImageView;
    }

    color_attachment_refs.resize(color_buffer_count);

    int color_buffer_index = 0;
    int attachment_point = 0;
    for (size_t i = 0; i < mAttachments.size(); ++i)
    {
        Texture2D *texture = mAttachments[i];
        if (texture->mImageAspectFlag == VK_IMAGE_ASPECT_COLOR_BIT)
        {
            attachments[i] = {};
            attachments[i].flags = 0;
            attachments[i].format = texture->mFormat;
            attachments[i].samples = VK_SAMPLE_COUNT_1_BIT;
            attachments[i].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
            attachments[i].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
            attachments[i].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
            attachments[i].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
            attachments[i].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            attachments[i].finalLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

            color_attachment_refs[color_buffer_index++] = {uint32_t(attachment_point++), VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL};
        }
        else if (texture->mImageAspectFlag == VK_IMAGE_ASPECT_DEPTH_BIT)
        {
            attachments[i] = {};
            attachments[i].flags = 0;
            attachments[i].format = mDepthBuffer->mFormat;
            attachments[i].samples = VK_SAMPLE_COUNT_1_BIT;
            attachments[i].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
            attachments[i].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
            attachments[i].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
            attachments[i].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
            attachments[i].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            attachments[i].finalLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

            depthAttachmentRef.attachment = attachment_point++;
            depthAttachmentRef.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
        }
        render_targets[i] = texture->mImageView;
    }

    VkSubpassDescription subpass{};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = color_attachment_refs.size();
    subpass.pColorAttachments = color_attachment_refs.data();
    subpass.pDepthStencilAttachment = &depthAttachmentRef;

    VkRenderPassCreateInfo rpci{};
    rpci.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    rpci.attachmentCount = attachments.size();
    rpci.pAttachments = attachments.data();
    rpci.subpassCount = 1;
    rpci.pSubpasses = &subpass;

    vkCreateRenderPass(xGetVulkanDevice(), &rpci, nullptr, &mRenderPass);

    VkFramebufferCreateInfo fbci{};
    fbci.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    fbci.pAttachments = render_targets.data();
    fbci.attachmentCount = render_targets.size();
    fbci.width = mWidth;
    fbci.height = mHeight;
    fbci.layers = 1;
    fbci.renderPass = mRenderPass;

    vkCreateFramebuffer(xGetVulkanDevice(), &fbci, nullptr, &mFBO);
}

void FrameBuffer::SetClearColor(int index, float r, float g, float b, float a)
{
    mClearValues[index].color = {r, g, b, a};
}

void FrameBuffer::SetClearDepthStencil(float depth, uint32_t stencil)
{
    mClearValues[mDepthBufferIndex].depthStencil = {depth, stencil};
}

VkCommandBuffer FrameBuffer::BeginRendering(VkCommandBuffer commandBuffer)
{
    VkCommandBuffer cmd;
    if (commandBuffer)
    {
        cmd = commandBuffer;
    }
    else
    {
        xBeginOneTimeCommandBuffer(&cmd);
    }

    VkFramebuffer render_target = mFBO;
    VkRenderPass render_pass = mRenderPass;

    VkRenderPassBeginInfo rpbi{};
    rpbi.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    rpbi.framebuffer = render_target;
    rpbi.renderPass = render_pass;
    rpbi.renderArea.offset = {0, 0};
    rpbi.renderArea.extent = {uint32_t(xGetViewportWidth()), uint32_t(xGetViewportHeight())};
    rpbi.clearValueCount = mClearValues.size();
    rpbi.pClearValues = mClearValues.data();

    vkCmdBeginRenderPass(cmd, &rpbi, VK_SUBPASS_CONTENTS_INLINE);

    return cmd;
}
